{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RoseTTAFold All-Atom (RFAA)本地部署\n",
    "\n",
    "## 相关链接\n",
    "\n",
    "- [Rohith Krishna et al. ,Generalized biomolecular modeling and design with RoseTTAFold All-Atom.Science384,eadl2528(2024).DOI:10.1126/science.adl2528](https://www.science.org/doi/10.1126/science.adl2528)\n",
    "- [github官方仓库](https://github.com/baker-laboratory/RoseTTAFold-All-Atom)\n",
    "- [SignalP-6.0h下载页面](https://services.healthtech.dtu.dk/cgi-bin/sw_request?software=signalp&version=6.0&packageversion=6.0h&platform=fast)\n",
    "\n",
    "## 开源协议:BSD\n",
    "\n",
    "## 介绍\n",
    "\n",
    "RoseTTAFold All-Atom (RFAA)基于RoseTTAFold2 蛋白结构预测模型,关注包含各种化合物的生物大分子组装体系结构预测.\n",
    "\n",
    "## RFAA本地部署\n",
    "\n",
    "### SignalP-6.0h下载\n",
    "\n",
    "SignalP-6.0h是预测信号肽位点的程序,需要填写学术许可协议.学术协议仅面向学术用户,商业用户需另行申请商业协议;用户不能在公共领域传播程序.\n",
    "\n",
    "打开[下载页面](https://services.healthtech.dtu.dk/cgi-bin/sw_request?software=signalp&version=6.0&packageversion=6.0h&platform=fast),填写姓名,职务,学术邮箱,隶属机构;阅读学术许可协议,接受协议\"I accept\"后提交\"Submit\".\n",
    "\n",
    "在学术邮箱中访问临时下载链接,可能被识别为垃圾邮件.下载链接目录下的`signalp-6.0h.fast.tar.gz`,约1.4G,推荐使用aria2,Motrix等工具多线程下载.\n",
    "\n",
    "### 环境配置\n",
    "\n",
    "在linux或mac操作系统下更新系统,打开终端,输入:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "# 不要在notebook中运行以下命令,仅用于方便复制\n",
    "\n",
    "# 安装miniforge,自带mamba包管理器,默认使用conda-forge通道\n",
    "# 可打开前半部分网址,下载aarch64架构或者mac操作系统的安装脚本\n",
    "# 使用conda进行包管理的用户不要重复安装,重复管理,以下使用mamba的命令用conda替换\n",
    "wget https://mirror.nju.edu.cn/github-release/conda-forge/miniforge/LatestRelease/Miniforge3-Linux-x86_64.sh\n",
    "chmod +x Miniforge3-Linux-x86_64.sh\n",
    "./Miniforge3-Linux-x86_64.sh\n",
    "# 按Enter 确定安装;按q 翻页;输入yes 同意协议;输入想要安装miniforge的路径,直接回车将默认安装在/home/regen/miniforge3;输入yes 初始化终端的mamba环境\n",
    "# mamba init fish # 安装时默认初始化bash,可指定初始化fish,可选\n",
    "bash # fish,重新打开终端,让修改生效\n",
    "\n",
    "# 设置conda和pip镜像\n",
    "pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple\n",
    "conda config --set show_channel_urls true\n",
    "conda config --set custom_channels.conda-forge https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/\n",
    "conda config --set custom_channels.bioconda https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/\n",
    "conda config --set custom_channels.pytorch https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/\n",
    "mamba env create -f environment.yaml -y # 创建RFAA环境\n",
    "mamba activate RFAA # 激活环境\n",
    "\n",
    "# git设置加速镜像,可选\n",
    "# git config --global url.\"https://gh-proxy.com/https://github.com/\".insteadOf https://github.com/\n",
    "\n",
    "git clone https://github.com/baker-laboratory/RoseTTAFold-All-Atom ~/git_develop/RFAA # 克隆仓库,仓库路径自行修改\n",
    "cd ~/git_develop/RFAA\n",
    "cd rf2aa/SE3Transformer/\n",
    "pip3 install --no-cache-dir -r requirements.txt\n",
    "python3 setup.py install\n",
    "cd -\n",
    "chmod +x input_prep/make_ss.sh\n",
    "signalp6-register ~/下载/signalp-6.0h.fast.tar.gz # 下载路径自行修改\n",
    "mv $CONDA_PREFIX/lib/python3.10/site-packages/signalp/model_weights/distilled_model_signalp6.pt $CONDA_PREFIX/lib/python3.10/site-packages/signalp/model_weights/ensemble_model_signalp6.pt # 兼容权重文件名\n",
    "bash install_dependencies.sh\n",
    "aria2c -x 16 -s 16 http://files.ipd.uw.edu/pub/RF-All-Atom/weights/RFAA_paper_weights.pt # 下载模型权重,wget 太慢\n",
    "aria2c -x 16 -s 16 https://ftp.ncbi.nlm.nih.gov/blast/executables/legacy.NOTSUPPORTED/2.2.26/blast-2.2.26-x64-linux.tar.gz # blast\n",
    "mkdir -p blast-2.2.26\n",
    "tar -xf blast-2.2.26-x64-linux.tar.gz -C blast-2.2.26\n",
    "cp -r blast-2.2.26/blast-2.2.26/ blast-2.2.26_bk\n",
    "rm -r blast-2.2.26\n",
    "mv blast-2.2.26_bk/ blast-2.2.26\n",
    "\n",
    "# 下载数据库\n",
    "# uniref30,46G,解压后181G\n",
    "aria2c -x 16 -s 16 http://wwwuser.gwdg.de/~compbiol/uniclust/2020_06/UniRef30_2020_06_hhsuite.tar.gz \n",
    "mkdir -p UniRef30_2020_06\n",
    "tar xfz UniRef30_2020_06_hhsuite.tar.gz -C ./UniRef30_2020_06\n",
    "# 结构模板,81G,解压后278G\n",
    "aria2c -x 16 -s 16 https://files.ipd.uw.edu/pub/RoseTTAFold/pdb100_2021Mar03.tar.gz\n",
    "tar xfz pdb100_2021Mar03.tar.gz\n",
    "# BFD,272G,解压后1.8T,如果uniref序列数量足够,将不会搜索该数据库\n",
    "# aria2c -x 16 -s 16 https://bfd.mmseqs.com/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz\n",
    "# mkdir -p bfd\n",
    "# tar xfz bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz -C ./bfd\n",
    "# 无法连接,使用af3的BFD镜像,数据库镜像地址为:\n",
    "# 没有下载过,不确定是否包含ffdata索引文件\n",
    "# aria2c -x 16 -s 16 https://storage.googleapis.com/alphafold-databases/casp14_versions/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz\n",
    "# mkdir -p bfd\n",
    "# tar --extract --verbose --file=bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt.tar.gz --directory=./bfd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 推理\n",
    "\n",
    "使用Hydra库处理预测设置.实际的推理脚本位于`rf2aa/run_inference.py`,默认设置文件为`rf2aa/config/inference/base.yaml`.用户自定义的配置文件也应该存放在该目录下.\n",
    "\n",
    "推荐使用默认参数,因为它们是训练时使用的参数,`loader_params.MAXCYCLE=10`除外,该参数默认为4,提高循环数量能够改善更困难的体系预测结果(论文中提及).\n",
    "\n",
    "进行推理的通用命令为:\n",
    "\n",
    "```py\n",
    "python -m rf2aa.run_inference --config-name 你的配置文件\n",
    "```\n",
    "\n",
    "主要的输入结构分为:\n",
    "\n",
    "- 蛋白 (protein_inputs)\n",
    "- 核酸 (na_inputs)\n",
    "- 小分子 (sm_inputs)\n",
    "- 蛋白和小分子之间的共价键\n",
    "- 修饰氨基酸,或者人造氨基酸(COMING SOON)\n",
    "\n",
    "以下是论文中不同预测任务的配置文件.\n",
    "\n",
    "#### 预测蛋白单体\n",
    "\n",
    "要求蛋白单体结构为fasta文件,可选的`job_name`用于输出文件名.示例配置文件为`rf2aa/config/inference/protein.yaml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from note_utils.path import expand_path,chdir\n",
    "\n",
    "RFAA_REPO = expand_path('~/git_develop/RFAA') # RFAA仓库路径自行修改"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "defaults:\n",
      "  - base\n",
      "\n",
      "job_name: \"7u7w_protein\"\n",
      "protein_inputs: \n",
      "  A:\n",
      "    fasta_file: examples/protein/7u7w_A.fasta\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !cat rf2aa/config/inference/protein.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "第1行表明继承base.yaml文件中的配置,其他参数定义后会覆盖默认配置.推荐修改`job_name`.\n",
    "\n",
    "`protein_inputs`中的`A`是链的id,必须存在,用于指定多条链.\n",
    "\n",
    "配置文件修改后执行推理:\n",
    "\n",
    "8G内存显卡显存不足,无法预测435 AA长度的蛋白.\n",
    "\n",
    "RFAA支持cpu推理,将pytorch的CUDA设备禁用即可,推理时间较长,包含7u7w的体系仅演示本例."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from contextlib import contextmanager\n",
    "\n",
    "@contextmanager\n",
    "def disable_gpu():\n",
    "    '''临时禁用GPU设备.'''\n",
    "    print('临时禁用GPU设备.')\n",
    "    os.environ['CUDA_VISIBLE_DEVICES']='\"\"'\n",
    "    yield\n",
    "    os.environ['CUDA_VISIBLE_DEVICES']='0'\n",
    "    print('重新启用GPU设备.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "临时禁用GPU设备.\n",
      "/home/regen/ana/envs/RFAA/lib/python3.10/site-packages/hydra/_internal/defaults_list.py:251: UserWarning: In 'protein': Defaults list is missing `_self_`. See https://hydra.cc/docs/1.2/upgrades/1.0_to_1.1/default_composition_order for more information\n",
      "  warnings.warn(msg, UserWarning)\n",
      "Using the cif atom ordering for TRP.\n",
      "重新启用GPU设备.\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO),disable_gpu():\n",
    "    !python -m rf2aa.run_inference --config-name protein"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SSD MSA搜索时长约20min,CPU推理时长约30min.预测的平均pae约为7.07."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "ATOM      1  N   GLY A   1      13.407 -33.989   5.700  1.00  0.62\n",
      "ATOM      2  CA  GLY A   1      12.108 -34.494   5.269  1.00  0.62\n",
      "ATOM      3  C   GLY A   1      10.999 -34.032   6.206  1.00  0.62\n",
      "ATOM      4  O   GLY A   1      10.540 -32.893   6.126  1.00  0.62\n",
      "ATOM      5  H   GLY A   1      14.238 -34.521   5.484  1.00  0.62\n",
      "\n",
      "预测的平均pae为:7.067476749420166\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !head -5 7u7w_protein.pdb\n",
    "    !echo\n",
    "    confidence= torch.load('7u7w_protein_aux.pt')\n",
    "    print(f'预测的平均pae为:{confidence[\"mean_pae\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 预测蛋白-核酸复合物\n",
    "\n",
    "包含额外的核酸链.以下是一个示例:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "defaults:\n",
      "  - base\n",
      "\n",
      "job_name: \"7u7w_protein_nucleic\"\n",
      "protein_inputs: \n",
      "  A: \n",
      "    fasta_file: examples/protein/7u7w_A.fasta\n",
      "na_inputs: \n",
      "  B: \n",
      "    fasta: examples/nucleic_acid/7u7w_B.fasta\n",
      "    input_type: \"dna\"\n",
      "  C: \n",
      "    fasta: examples/nucleic_acid/7u7w_C.fasta\n",
      "    input_type: \"dna\"\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !cat rf2aa/config/inference/nucleic_acid.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "双股螺旋DNA必须明确使用两条链表示.\n",
    "\n",
    "核酸的允许输入类型包括dna和rna.\n",
    "\n",
    "目前不支持RNA的MSA,或者将蛋白的MSA与RNA的MSA配对.目前请使用RF-NA模型进行结构预测.\n",
    "\n",
    "预测复合物的命令如下:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with chdir(RFAA_REPO),disable_gpu():\n",
    "    !python -m rf2aa.run_inference --config-name nucleic_acid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 预测蛋白-小分子复合物\n",
    "\n",
    "小分子输入格式可以是`smiles`字符串,或者openbabel支持读取的分子结构格式(sdf,smi,mol2,pdb等).其中小分子的`input`和`input_type`必须同时提供.金属离子的输入格式需要是`sdf`文件或者`smiles`字符串.\n",
    "\n",
    "配置文件和推理命令示例如下:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "defaults:\n",
      "  - base\n",
      "\n",
      "job_name: 7qxr\n",
      "\n",
      "protein_inputs:\n",
      "  A: \n",
      "    fasta_file: examples/protein/7qxr.fasta\n",
      "\n",
      "sm_inputs:\n",
      "  B:\n",
      "    input: examples/small_molecule/NSW_ideal.sdf\n",
      "    input_type: \"sdf\"\n",
      "/home/regen/ana/envs/RFAA/lib/python3.10/site-packages/hydra/_internal/defaults_list.py:251: UserWarning: In 'protein_sm': Defaults list is missing `_self_`. See https://hydra.cc/docs/1.2/upgrades/1.0_to_1.1/default_composition_order for more information\n",
      "  warnings.warn(msg, UserWarning)\n",
      "Using the cif atom ordering for TRP.\n",
      "./make_msa.sh examples/protein/7qxr.fasta 7qxr/A 4 64  pdb100_2021Mar03/pdb100_2021Mar03\n",
      "Predicting: 100%|██████████████████████████| 1/1 [00:00<00:00,  1.95sequences/s]\n",
      "Running HHblits against UniRef30 with E-value cutoff 1e-10\n",
      "- 16:51:37.275 INFO: Input file = 7qxr/A/hhblits/t000_.1e-10.a3m\n",
      "\n",
      "- 16:51:37.275 INFO: Output file = 7qxr/A/hhblits/t000_.1e-10.id90cov75.a3m\n",
      "\n",
      "- 16:51:37.373 WARNING: Maximum number 100000 of sequences exceeded in file 7qxr/A/hhblits/t000_.1e-10.a3m\n",
      "\n",
      "- 16:53:28.773 INFO: Input file = 7qxr/A/hhblits/t000_.1e-10.a3m\n",
      "\n",
      "- 16:53:28.774 INFO: Output file = 7qxr/A/hhblits/t000_.1e-10.id90cov50.a3m\n",
      "\n",
      "- 16:53:28.870 WARNING: Maximum number 100000 of sequences exceeded in file 7qxr/A/hhblits/t000_.1e-10.a3m\n",
      "\n",
      "Running PSIPRED\n",
      "Running hhsearch\n",
      "cat: 7qxr/A/t000_.ss2: 没有那个文件或目录\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !cat rf2aa/config/inference/protein_sm.yaml\n",
    "    !echo \n",
    "    !python -m rf2aa.run_inference --config-name protein_sm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其中`cat: 7qxr/A/t000_.ss2: 没有那个文件或目录`报错是因为UniRef30搜索过程中已经搜索了足够的序列,bfd数据库跳过了,对推理没有影响.\n",
    "\n",
    "GPU推理时长约30s.\n",
    "\n",
    "由于输出路径默认为\"\",输出的pdb和指标文件将生成在RFAA仓库根目录下."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "ATOM      1  N   THR A   1      -4.772  21.530  -3.057  1.00  0.55\n",
      "ATOM      2  CA  THR A   1      -4.991  20.092  -2.955  1.00  0.55\n",
      "ATOM      3  C   THR A   1      -5.965  19.606  -4.021  1.00  0.55\n",
      "ATOM      4  O   THR A   1      -7.004  20.224  -4.254  1.00  0.55\n",
      "ATOM      5  CB  THR A   1      -5.555  19.765  -1.560  1.00  0.55\n",
      "\n",
      "9.76700210571289\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !head -5 7qxr.pdb\n",
    "    !echo\n",
    "    confidence= torch.load('7qxr_aux.pt')\n",
    "    print(f'预测的分子间pae为:{confidence[\"pae_inter\"]},<10表示高的对接质量')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 预测更复杂的复合物体系\n",
    "\n",
    "以下是包含蛋白,核酸,小分子的体系:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with chdir(RFAA_REPO),disable_gpu():\n",
    "    !cat rf2aa/config/inference/protein_na_sm.yaml\n",
    "    !echo \n",
    "    !python -m rf2aa.run_inference --config-name protein_na_sm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 预测共价修饰蛋白\n",
    "\n",
    "注意:指定共价修饰会引入新问题:\n",
    "\n",
    "- 形成或者消除手性中心.RFAA需要指定手性信息.手性中心由openbabel识别,但结果不总是和化学直觉相符\n",
    "- 共价修饰总是有\"离去基团\",在修饰时离开蛋白和修饰化学物质\n",
    "\n",
    "要定义共价键,使用以下信息定义:\n",
    "\n",
    "```py\n",
    "(protein_chain, residue_number, atom_name), (small_molecule_chain, atom_index), (new_chirality_atom_1, new_chirality_atom_2)\n",
    "```\n",
    "\n",
    "蛋白残基编号和原子索引都从1开始.多数情况下原子的手性不会变化,那么手性为(\"null\", \"null\").\n",
    "\n",
    "可选的手性是ccw和cw,分别对应逆时针counterclockwise S和顺时针clockwise,R型.\n",
    "\n",
    "如果用户未指定手性,openbabel会在检测到新的手性中心时报错.即使你确定openbabel的检测是错的,在训练过程中也是以openbabel的手性信息训练的.推荐你指定这些位置的手性,以获得最好的效果.\n",
    "\n",
    "你不能定义两个小分子之间的共价键,如果有些PDB文件中将小分子定义为多个残基,你必须合并为1个残基.\n",
    "\n",
    "你必须在输入分子中删除离去基团.蛋白侧链的离去基团会自动修改.目前hydra配置文件中允许动态提供离去基团,但只是实验功能,未经完整测试.\n",
    "\n",
    "以下是一个示例:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with chdir(RFAA_REPO):\n",
    "    !cat rf2aa/config/inference/covalent.yaml\n",
    "    !echo \n",
    "    !python -m rf2aa.run_inference --config-name covalent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "序列不足,需要搜索bfd,不作演示."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "共价修饰的小分子必须是sdf格式.原始的输入为:\n",
    "\n",
    "[((\"A\", \"74\", \"ND2\"), (\"B\", \"1\"), (\"CW\", \"null\"))].\n",
    "\n",
    "为了处理hydra的输入要求,共价键输入中间的双引号被转义:\n",
    "\n",
    "`\"[((\\\"A\\\", \\\"74\\\", \\\"ND2\\\"), (\\\"B\\\", \\\"1\\\"), (\\\"CW\\\", \\\"null\\\"))]\"`\n",
    "\n",
    "#### 理解模型输出\n",
    "\n",
    "输出结果包括:\n",
    "\n",
    "- PDB预测结构文件,温度因子表示每个位置的plddt\n",
    "- 存储置信指标的pytorch文件,可用`torch.load(file, map_location=\"cpu\")`加载.\n",
    "\n",
    "置信指标包括:\n",
    "\n",
    "1. plddts,节点级别的plddt张量\n",
    "2. pae,LxL的张量,预测第i个位置的框架frame对齐后每个位置的对齐误差(对于原子节点则是原子框架)\n",
    "3. pde,LxL的张量,预测每个成对距离的无符号误差\n",
    "4. mean_plddt,plddts平均值\n",
    "5. mean_pae,pae的平均值\n",
    "6. pae_prot,蛋白残基级别的pae平均值\n",
    "7. pae_inter,蛋白残基与原子框架之间的误差,以及原子坐标与蛋白框架之间的误差均值.pae_inter<10表示高质量的对接结果"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RFAA 个人测试用例:SEA蛋白金属离子结合位点预测\n",
    "\n",
    "pdbid:1esf,为金葡菌肠毒素A,锌离子结合位点是与MHC-II结合并发挥作用的重要位点.实验使用同族的镉离子.\n",
    "\n",
    "AF2,AF3对实验缺失的N端序列预测不合理,本实验主要测试RFAA对缺失N端的预测效果,以及较少示例的镉离子位点预测效果.\n",
    "\n",
    "需要搜索bfd,临时修改了uniref数据库接受序列的下限(make_msa.sh中n75>2000改为>50),跳过bfd数据库."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "defaults:\n",
      "  - base\n",
      "job_name: 1esf\n",
      "protein_inputs:\n",
      "  A: \n",
      "    fasta_file: 1esf/P0A0L2.fasta\n",
      "sm_inputs:\n",
      "  B:\n",
      "    input: \"[Cd+2]\"\n",
      "    input_type: \"smiles\"\n",
      "/home/regen/ana/envs/RFAA/lib/python3.10/site-packages/hydra/_internal/defaults_list.py:251: UserWarning: In '1esf.rfaa': Defaults list is missing `_self_`. See https://hydra.cc/docs/1.2/upgrades/1.0_to_1.1/default_composition_order for more information\n",
      "  warnings.warn(msg, UserWarning)\n",
      "Using the cif atom ordering for TRP.\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO) as (tmp,origin):\n",
    "    !cat $origin/1esf/1esf.rfaa.yaml\n",
    "    !cp -R $origin/1esf/ $tmp\n",
    "    !cp 1esf/1esf.rfaa.yaml rf2aa/config/inference/1esf.rfaa.yaml # 仅支持在rf2aa/config/inference/中搜索配置文件,除非改源代码\n",
    "    !echo\n",
    "    !python -m rf2aa.run_inference --config-name 1esf.rfaa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "当前目录从 /home/regen/git_develop/rosettafold_notebook 临时切换到 /home/regen/git_develop/RFAA\n",
      "预测的分子间pae为:17.68967056274414,<10表示高的对接质量\n",
      "预测的蛋白pae为:4.212173938751221\n",
      "当前目录切换回到 /home/regen/git_develop/rosettafold_notebook\n"
     ]
    }
   ],
   "source": [
    "with chdir(RFAA_REPO) as (tmp,origin):\n",
    "    confidence= torch.load('1esf_aux.pt')\n",
    "    print(f'预测的分子间pae为:{confidence[\"pae_inter\"]},<10表示高的对接质量')\n",
    "    print(f'预测的蛋白pae为:{confidence[\"pae_prot\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RFAA环境依赖与pymol有冲突,生成pml文件,在其他环境打开pymol:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp $RFAA_REPO/1esf.pdb 1esf/pred.pdb\n",
    "with open('1esf/1esf.pml','w') as pml:\n",
    "    cmds='''\n",
    "load 1esf/1esf.cif,native\n",
    "color grey,native\n",
    "load 1esf/pred.pdb,pred\n",
    "color blue,pred\n",
    "color green,pred and resn CD\n",
    "util.cnc\n",
    "remove solvent or (native and chain B)\n",
    "super pred,native\n",
    "\n",
    "'''\n",
    "    pml.write(cmds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在有pymol的环境下,终端输入`pymol 1esf/1esf.pml`,RMSD约0.6,蛋白预测效果可接受,N端差异较明显.\n",
    "\n",
    "fasta文件中提供了N端缺失的序列,但在预测结果中直接被截断,并未建模.可能原因是置信指标太低,拒绝输出.\n",
    "\n",
    "Cd离子未正确识别,被建模为LG1未知原子.但结合位点预测正确.\n",
    "\n",
    "![](1esf/pred_with_seq.png)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "RFAA",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
